﻿using System;
using System.Diagnostics;
using CoreCalc;
using RayDen.Library.Components.Surface;
using RayDen.Library.Core;
using RayDen.Library.Core.Primitives;
using RayDen.Library.Data.Imaging;
using RayDen.RayEngine.Core.Interface;
using RayDen.RayEngine.Core.Types;
using RayDen.RayEngine.Engines.PathTracer;
using RayDen.RayEngine.Samplers;

namespace RayDen.RayEngine.Engines.SpectralPT
{

    public class UnidirectionalSpectralPathSampler : PathSamplerBase
    {
        internal SampledSpectrum Throughput;
        public PathTracerPathState PathState;



        protected int depth, tracedShadowRayCount;
        protected float pathWeight;
        protected bool specularBounce, inside;

        public UInt64 PathExpression;
        public new SampledSpectrum Radiance;

        protected RayEngineScene scene;
        private SpectralShadowRayInfo[] secRays;
        SurfaceIntersectionData hitInfo = null;
        public int MaxRaysPerPath
        {
            get;
            set;
        }


        public override void InitPath(IPathProcessor buffer)
        {
            base.InitPath(buffer);
            this.scene = pathIntegrator.Scene;
            this.Radiance = new SampledSpectrum(0f);
            this.Throughput = new SampledSpectrum(1f);
            
            this.PathState = PathTracerPathState.EyeVertex;
            if (this.secRays == null)
                this.secRays = new SpectralShadowRayInfo[scene.ShadowRaysPerSample];
            if (this.Sample == null)
                this.Sample = pathIntegrator.Sampler.GetSample(null);
            else
                pathIntegrator.Sampler.GetSample(this.Sample);
            pathIntegrator.Scene.Camera.GetRay(Sample.imageX, Sample.imageY, out this.PathRay);
            this.RayIndex = -1;
            //this.pathWeight = 1.0f;
            this.tracedShadowRayCount = 0;
            this.depth = 0;
            this.specularBounce = true;
            inside = false;
        }


        public override void Splat(SampleBuffer sampleBuffer)
        {

            ColorFactory.ToRgb(ref this.Radiance, ref base.Radiance);//this.Radiance.ToRgb();
            base.Splat(sampleBuffer);
        }

        public override bool FillRayBuffer(RayBuffer rayBuffer)
        {
            var leftSpace = rayBuffer.LeftSpace();
            if (((PathState == PathTracerPathState.EyeVertex) && (1 > leftSpace)) ||
            ((PathState == PathTracerPathState.ShadowRaysOnly) && (tracedShadowRayCount > leftSpace)) ||
            ((PathState == PathTracerPathState.NextVertex) && (tracedShadowRayCount + 1 > leftSpace)))
                return false;
            if (PathState != PathTracerPathState.ShadowRaysOnly)
                RayIndex = rayBuffer.AddRay(ref PathRay);
            if (PathState == PathTracerPathState.NextVertex || PathState == PathTracerPathState.ShadowRaysOnly)
            {
                for (int i = 0; i < tracedShadowRayCount; ++i)
                    secRays[i].currentShadowRayIndex = rayBuffer.AddRay(ref secRays[i].shadowRay);
            }
            return true;
        }


        public override void Advance(RayBuffer rayBuffer, SampleBuffer consumer)
        {
            int currentTriangleIndex = 0;
#if VERBOSE

            try {
#endif

            base.Advance(rayBuffer, consumer);

            RayHit rayHit = rayBuffer.rayHits[RayIndex];

            if (((PathState == PathTracerPathState.ShadowRaysOnly) || (PathState == PathTracerPathState.NextVertex)) &&
                (tracedShadowRayCount > 0))
            {
                for (int i = 0; i < tracedShadowRayCount; ++i)
                {
                    RayHit shadowRayHit = rayBuffer.rayHits[secRays[i].currentShadowRayIndex];
                    RgbSpectrum attenuation;
                    bool continueTrace;
                    if (this.ShadowRayTest(ref shadowRayHit, ref secRays[i].shadowRay, out attenuation, out continueTrace))
                    {
                        //                            Radiance.MADD()
                        //Radiance.MAdd(ref secRays[i].color, 1f / secRays[i].pdf);
                        SSE.MAddSSE(ref secRays[i].color, 1f / secRays[i].pdf, ref Radiance);
                        //Radiance += ((secRays[i].color) / secRays[i].pdf);
                        //pathWeight *= secRays[i].pdf;
                    }
                }
                if (PathState == PathTracerPathState.ShadowRaysOnly)
                {
                    Splat(consumer);
                    return;
                }
                tracedShadowRayCount = 0;

            }

            depth++;
            bool missed = rayHit.Index == 0xffffffffu;
            if (missed || PathState == PathTracerPathState.ShadowRaysOnly || depth > scene.MaxPathDepth)
            {
                if (missed)
                {
                    SampledSpectrum env;
                    this.SampleEnvironment(ref PathRay.Dir, out env);
                    //Radiance.MAdd(ref env, ref Throughput);
                    SSE.MAddSSE(ref env, ref Throughput, ref Radiance);
                    //Radiance += env * Throughput;
                }
                Splat(consumer);

                return;
            }

            // Something was hit
            if (hitInfo == null)
            {
                hitInfo = SurfaceSampler.GetIntersection(ref PathRay, ref rayHit, IntersectionOptions.ResolveSurfaceComplete | IntersectionOptions.ResolveSpectralDistributions);
            }
            else
            {
                SurfaceSampler.GetIntersection(ref PathRay, ref rayHit, ref hitInfo, IntersectionOptions.ResolveSurfaceComplete | IntersectionOptions.ResolveSpectralDistributions);
            }
            currentTriangleIndex = (int)rayHit.Index;

            if (hitInfo == null)
            {
                Debugger.Break();
            }
            Vector wo = -PathRay.Dir;
            RgbSpectrum rThroughput = RgbSpectrum.UnitSpectrum();
            //If Hit light)
            if (hitInfo.IsLight)
            {
                if (specularBounce || depth == 1)
                {
                    var lt = scene.GetLightByIndex(currentTriangleIndex);
                    if (lt != null)
                    {
                        var le = (SampledSpectrum)lt.Le(ref wo); 
                        //Radiance += Throughput * le;
                        //Radiance.MAdd(ref le, ref Throughput);
                        SSE.MAddSSE(ref le, ref Throughput, ref Radiance);

                    }
                }
                Splat(consumer);

                return;
            }

            var hitPoint = PathRay.Point(rayHit.Distance);

            var bsdf = hitInfo.MMaterial;

            if (bsdf.IsDiffuse())
            {
                float lightStrategyPdf = LightSampler.StrategyPdf;
                //scene.ShadowRaysPerSample/(float)scene.Lights.Length;

                SampledSpectrum lightTroughtput = hitInfo.Reflectance;
                SSE.MulSSE(ref Throughput, ref lightTroughtput, ref lightTroughtput);

                LightSampler.EvaluateShadow(ref hitPoint, ref hitInfo.Normal, Sample.GetLazyValue(),
                                            Sample.GetLazyValue(), Sample.GetLazyValue(), ref ls);
                foreach (var lightSample in ls)
                {
                    if (lightSample.Pdf <= 0f)
                        continue;
                    secRays[tracedShadowRayCount].color = (SampledSpectrum) lightSample.Spectrum;
                    secRays[tracedShadowRayCount].pdf = lightSample.Pdf;
                    secRays[tracedShadowRayCount].shadowRay = lightSample.LightRay;
                    Vector lwi = secRays[tracedShadowRayCount].shadowRay.Dir;
                    RgbSpectrum fs;
                    ColorFactory.ToRgb(ref this.Throughput, ref rThroughput);
                        //Throughput.ToRgb();

                    hitInfo.MMaterial.f(
                        ref secRays[tracedShadowRayCount].shadowRay.Dir,
                        ref wo, ref hitInfo.ShadingNormal, ref rThroughput, out fs, types: BrdfType.Diffuse);

                    SampledSpectrum tmp;
                    SampledSpectrum ft = ColorFactory.FromRgb(ref fs, SpectrumType.Reflectance);
                    SSE.MulSSE(ref lightTroughtput, Vector.AbsDot(ref hitInfo.Normal, ref lwi), ref tmp);
                    SSE.MulSSE(ref tmp, ref ft, ref tmp);
                    SSE.MulSSE(ref secRays[tracedShadowRayCount].color, ref tmp, ref secRays[tracedShadowRayCount].color);
                    /*
                    secRays[tracedShadowRayCount].color *= lightTroughtput *
                                                           Vector.AbsDot(ref hitInfo.Normal, ref lwi) *
                                                           ColorFactory.FromRgb(ref fs, SpectrumType.Reflectance);*/

                    if (!secRays[tracedShadowRayCount].color.IsBlack())
                    {
                        secRays[tracedShadowRayCount].pdf /= lightStrategyPdf;
                        tracedShadowRayCount++;
                    }
                }
            }

            float fPdf = 0f;
            var wi = new Vector();
            RgbSpectrum f;
            ColorFactory.ToRgb(ref this.Throughput, ref rThroughput);
                //Throughput.ToRgb();
            BsdfEvent evt;
            f = bsdf.Sample_f(ref wo, out wi, ref hitInfo.Normal, ref hitInfo.ShadingNormal, ref rThroughput,
                                Sample.GetLazyValue(), Sample.GetLazyValue(),
                                Sample.GetLazyValue(), ref hitInfo.TextureData,
                                out fPdf, out evt) ;
            specularBounce = evt.Has(BsdfEvent.Specular);

            if ((fPdf <= 0.0f) || f.IsBlack())
            {
                if (tracedShadowRayCount > 0)
                    PathState = PathTracerPathState.ShadowRaysOnly;
                else
                {
                    Splat(consumer);

                }
                return;
            }
            

            var ff = ColorFactory.FromRgb(ref f);
            SSE.MulSSE(ref ff, ref hitInfo.Reflectance, ref ff);

            //if (inside) { fPdf /= MathLab.Exp(-rayHit.Distance*ff.y());}

            SSE.MulSSE(ref ff, ref Throughput, ref Throughput);
            SSE.MulSSE(ref Throughput, 1f / fPdf, ref Throughput);

            //Throughput *= ( * hitInfo.Reflectance) / fPdf;
            inside = evt.Has(BsdfEvent.Transmit);


            if (depth > scene.MaxPathDepth)
            {
                float prob = Math.Max(Throughput.Filter(), scene.RussianRuletteImportanceCap);
                if (prob >= Sample.GetLazyValue())
                {
                    //Throughput /= prob;
                    SSE.MulSSE(ref Throughput, 1f / prob, ref Throughput);
                    //pathWeight *= prob;
                }
                else
                {
                    if (tracedShadowRayCount > 0)
                        PathState = PathTracerPathState.ShadowRaysOnly;
                    else
                    {
                        Splat(consumer);

                    }

                    return;
                }
            }

            PathRay.Org = hitPoint;
            PathRay.Dir = wi.Normalize();
            PathState = PathTracerPathState.NextVertex;

#if VERBOSE
            }
            catch (Exception ex) {
                RayDen.Library.Components.SystemComponents.Tracer.TraceLine("Advance path exception");
                RayDen.Library.Components.SystemComponents.Tracer.TraceLine("Error triangle {0}", currentTriangleIndex);
                RayDen.Library.Components.SystemComponents.Tracer.TraceLine(ex.Message);
                RayDen.Library.Components.SystemComponents.Tracer.TraceLine(ex.StackTrace);

            }
#endif
        }

        protected override bool ShadowRayTest(ref RayHit shadowRayHit, ref RayData shadowRay, out RgbSpectrum attenuation, out bool continueTrace)
        {
            var hit = shadowRayHit.Index == 0xffffffffu || scene.IsLight((int)shadowRayHit.Index);
            attenuation = new RgbSpectrum(1f);
            continueTrace = false;

            return hit;
        }

        private void SampleEnvironment(ref Vector vector, out SampledSpectrum result)
        {
            result = ColorFactory.FromRgb(this.LightSampler.SampleEnvironment(-vector), SpectrumType.Illuminant);
        }

    }
}
